from UCF.FileLoader import FileLoader
from typing import Set, List, Tuple, Dict
from tqdm import tqdm
import pickle
import os


class Trainer:
    def __init__(self, hyper_params):
        self.file_loader = FileLoader(hyper_params)

        # Create Interaction Set
        print('UCF: Trainer: Loading Interaction List...')
        self.interaction_list: Dict[int, List[int]] = {}
        self.type_list: Dict[int, List[int]] = {}

        for user, history in self.file_loader.intereaction_list.items():
            self.interaction_list[user] = []
            self.type_list[user] = []
            for item in history:
                self.interaction_list[user].append(item.item_id)
                self.type_list[user].append(item.type_id)
        print('UCF: Trainer: Finished.')

        try:
            self.W = pickle.load(open('UCF/model_dat/W.pkl', 'rb'))
            print('UCF: Model Loaded.')
        except Exception as err:
            self.cal_similarity()

    def cal_similarity(self):
        # Build inverse table
        item_user: Dict[int, Set[int]] = {}
        for u, items in self.type_list.items():
            for i in items:
                if i not in item_user.keys():
                    item_user[i] = set()
                item_user[i].add(u)
        print('UCF: Inverse table finished')

        # Co-rated items between users
        C = {}
        N = {}
        W = {}

        print('UCF: Initializing...')
        for u in self.type_list.keys():
            N[u] = 0
            C[u] = {}
            W[u] = {}

        print('UCF: Counting...')
        for i, users in tqdm(item_user.items()):
            for u in users:
                N[u] += 1
                for v in users:
                    if v == u:
                        continue
                    if v not in C[u]:
                        C[u][v] = 0
                    C[u][v] += 1

        print('UCF: Co-rated items count finished')
        # Calculate similarity matrix
        for u, related_users in tqdm(C.items()):
            for v, cuv in related_users.items():
                W[u][v] = cuv/(N[u]*N[v])**0.5
        print('UCF: Similarity calculation finished')

        self.W = W
        os.makedirs('UCF/model_dat', exist_ok=True)
        pickle.dump(self.W, open('UCF/model_dat/W.pkl', 'wb'))


if __name__ == '__main__':
    hyper_params = {
        'dataset_path': '../datasets/lib.txt'
    }

    trainer = Trainer(hyper_params)
